package com.mashibing.common.algorithm;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import com.mashibing.common.exception.RRException;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 树节点工具类
 *
 * @author ROGzhongyang
 * @date 2024/11/28 06:56:56
 */
public class TreeNodeUtils {

    /**
     * 顶级父 id ,默认为 0
     */
    private static final Long TOP_PARENT_ID = 0L;

    /**
     * 使用 BFS 算法 返回树节点列表
     * @param treeNode 原始节点列表
     * @return 填充后的树节点列表
     */
    public static <T> List<TreeNode<T>> getTreeListByBFS(List<TreeNode<T>> treeNode){

        // 获取所有父 id 的集合
        Set<Long> parentIdSet = treeNode.stream().distinct().map(TreeNode::getParentId).collect(Collectors.toSet());

        // 创建一个映射，用于存储父节点 id 和对应的子节点列表的关系
        Map<Long, List<TreeNode<T>>> parentToChildrenMap = treeNode.stream().collect(Collectors.groupingBy(TreeNode::getParentId));

        // 清空树节点列表，以便后续复用
        treeNode.clear();

        // 创建一个队列用于 BFS 遍历
        Queue<TreeNode<T>> queue = new LinkedList<>();

        // 找到顶级节点并放入队列和树节点列表
        for (TreeNode<T> node: parentToChildrenMap.get(TOP_PARENT_ID)) {
            if (ObjectUtil.isEmpty(node))
                throw new RRException("顶级节点的父节点不为" + TOP_PARENT_ID + "!", 20);
            queue.offer(node);
            treeNode.add(node);
        }

        // 当队列不为空时，进行 BFS 遍历
        while (CollectionUtil.isNotEmpty(queue)) {

            // 获取队列头部节点
            TreeNode<T> current = queue.poll();

            if (current != null) {

                // 如果当前节点 id 包含在父 id 集合内
                if (parentIdSet.contains(current.getId())) {

                    // 根据映射设置当前节点的子节点列表
                    List<TreeNode<T>> nodeList = parentToChildrenMap.getOrDefault(current.getId(), Collections.emptyList());
                    current.setChildren(nodeList);

                    // 将当前节点的子节点列表放入队列，以便后续处理
                    queue.addAll(nodeList);

                }

                // 如果当前节点是叶子节点，则将其children置为空
                else
                    current.setChildren(null);
            }
        }
        return treeNode;
    }
}
